#!/usr/bin/env python3
# -*-coding:utf8-*-
# 通过游戏手柄控制机械臂末端位姿

from re import X
from typing import Optional, Tuple, Dict, List
import time
import math
import pygame
from piper_sdk import *

# 设置转换因子，用于将物理单位转换为机械臂控制单位
factor = 1000
HOME = [55.0, 0.0, 206.0, 0, 85.0, 0, 0]

class JoystickController:
    """
    游戏手柄控制器类，用于读取手柄输入并生成六自由度速度
    """
    def __init__(self, joystick_id: int = 0,piper=None):
        """
        初始化游戏手柄控制器
        
        参数:
            joystick_id: 游戏手柄ID，默认为0（第一个连接的手柄）
        """
        self.piper = piper
        # 初始化pygame和joystick模块
        pygame.init()
        pygame.joystick.init()
        
        # 检查是否有可用的游戏手柄
        self.joystick_count = pygame.joystick.get_count()
        if self.joystick_count == 0:
            raise Exception("未检测到游戏手柄，请连接游戏手柄后重试")
        
        # 选择指定ID的游戏手柄
        if joystick_id >= self.joystick_count:
            print(f"警告: 指定的游戏手柄ID {joystick_id} 不可用，使用ID 0")
            joystick_id = 0
            
        # 初始化游戏手柄
        self.joystick = pygame.joystick.Joystick(joystick_id)
        self.joystick.init()
        
        # 打印游戏手柄信息
        print(f"已连接游戏手柄: {self.joystick.get_name()}")
        print(f"轴数量: {self.joystick.get_numaxes()}")
        print(f"按钮数量: {self.joystick.get_numbuttons()}")
        print(f"帽子开关数量: {self.joystick.get_numhats()}")
        print(f"帽子开关数量: {self.joystick.get_numhats()}")
        
        # 初始化速度参数
        self.speed_factor = 0.3  # 速度因子，用于调整速度大小
        self.linear_speed_max = 0.1  # 最大线速度 (米/秒)
        self.angular_speed_max = 3.0  # 最大角速度 (度/秒)
        
        # 初始化六自由度速度
        self.vx = 0.0  # X轴线速度 (米/秒)
        self.vy = 0.0  # Y轴线速度 (米/秒)
        self.vz = 0.0  # Z轴线速度 (米/秒)
        self.rx = 0.0  # X轴角速度 (度/秒)
        self.ry = 0.0  # Y轴角速度 (度/秒)
        self.rz = 0.0  # Z轴角速度 (度/秒)
        
        # 定义按钮映射（根据实际手柄可能需要调整）
        self.button_map = {
            'y': 2,
            'x': 0,
            'speed_up': 4,    # 通常是LB按钮
            'speed_down': 5,  # 通常是RB按钮
            'reset': 9,       # 通常是Start按钮
            'stop': 8,        # 通常是Back/Select按钮
        }
        
        # 定义轴映射（根据实际手柄可能需要调整）
        self.axis_map = {
            'left_x': 0,      # 左摇杆X轴
            'left_y': 1,      # 左摇杆Y轴
            'right_z': 3,     # 右摇杆X轴 R轴移动
            'right_y': 4,     # 右摇杆Y轴 rz角速度
            'lt': 2,          # 左触发键
            'rt': 5,          # 右触发键
            'rx': 6,          # X轴角速度
            'ry': 7,          # Y轴角速度
        }
        
        # 设置死区值，避免摇杆轻微偏移导致意外移动
        self.deadzone = 0.15

    def get_key_name(self, key: int) -> str:
        for k, v in self.button_map.items():
            if v == key:
                return k
            # print(f"key:{k}, unknown value:{v}")
        return "unknown"
    def process_events(self) -> bool:
        """
        处理pygame事件
        
        返回:
            bool: 如果应该退出程序则返回False，否则返回True
        """
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            
            # 处理按钮按下事件
            elif event.type == pygame.JOYBUTTONDOWN:
                print(f"按钮按下: {event.button} key:{self.get_key_name(event.button)}")
                if event.button == self.button_map['speed_up']:
                    self.speed_factor = min(self.speed_factor + 0.1, 1.0)
                    print(f"速度因子增加至: {self.speed_factor:.1f}")
                elif event.button == self.button_map['speed_down']:
                    self.speed_factor = max(self.speed_factor - 0.1, 0.1)
                    print(f"速度因子减小至: {self.speed_factor:.1f}")
                elif event.button == self.button_map['reset']:
                    X = round(HOME[0]*factor)
                    Y = round(HOME[1]*factor)
                    Z = round(HOME[2]*factor)
                    RX = round(HOME[3]*factor)
                    RY = round(HOME[4]*factor)
                    RZ = round(HOME[5]*factor)
                    joint_6 = round(HOME[6]*factor)
                    print(X,Y,Z,RX,RY,RZ)
                    self.piper.MotionCtrl_2(0x01, 0x00, 100, 0x00)
                    self.piper.EndPoseCtrl(X,Y,Z,RX,RY,RZ)
                    self.piper.GripperCtrl(abs(joint_6), 1000, 0x01, 0)
                elif event.button == self.button_map['stop']:
                    print("紧急停止")
                    self.piper.MotionCtrl_1(0x01,0,0)
        
        return True
    
    
    def apply_deadzone(self, value: float) -> float:
        """
        应用死区值到摇杆输入
        
        参数:
            value: 摇杆输入值 (-1.0 到 1.0)
            
        返回:
            应用死区后的值
        """
        if abs(value) < self.deadzone:
            return 0.0
        
        # 重新映射值，使其在死区外从0到1平滑过渡
        sign = 1.0 if value > 0 else -1.0
        return sign * (abs(value) - self.deadzone) / (1.0 - self.deadzone)
    
    def update_velocity(self):
        """
        根据游戏手柄输入更新六自由度速度
        """
        # self.joystick.get_hat(0) 返回(0,0)
        griper = -self.joystick.get_hat(0)[0]
        if griper == 1:
            self.piper.GripperCtrl(round(0.08*1000*1000), 1000, 0x01, 0)
        elif griper == -1:
            self.piper.GripperCtrl(0, 1000, 0x01, 0)
        # 读取左摇杆输入 (用于控制X和Y轴线速度)
        left_x = self.apply_deadzone(self.joystick.get_axis(self.axis_map['left_x']))
        left_y = self.apply_deadzone(self.joystick.get_axis(self.axis_map['left_y']))
        
        # 读取右摇杆输入 (用于控制Z轴线速度和RZ轴角速度)
        # right_x = self.apply_deadzone(self.joystick.get_axis(self.axis_map['right_x']))
        right_y = self.apply_deadzone(self.joystick.get_axis(self.axis_map['right_y']))
        
        # 读取触发键输入 (用于控制RX和RY轴角速度)
        lt = (self.joystick.get_axis(self.axis_map['lt']) + 1.0) / 2.0  # 将-1到1映射为0到1
        rt = (self.joystick.get_axis(self.axis_map['rt']) + 1.0) / 2.0  # 将-1到1映射为0到1
        
        # 更新线速度 (毫米/秒)
        self.vx = -left_y * self.linear_speed_max * self.speed_factor * factor  # 前后移动
        self.vy = -left_x * self.linear_speed_max * self.speed_factor * factor  # 左右移动
        self.vz = -right_y * self.linear_speed_max * self.speed_factor * factor  # 上下移动
        
        # 更新角速度 (度/秒)
        # self.rx = self.apply_deadzone(self.joystick.get_axis(self.axis_map['rx']))   # 绕X轴旋转
        self.ry = (lt - rt) * self.angular_speed_max * self.speed_factor    # 绕Y轴旋转
        self.rz = -self.apply_deadzone(self.joystick.get_axis(self.axis_map['right_z']))   # 绕Z轴旋转
    
    def get_velocity(self) -> Tuple[float, float, float, float, float, float]:
        """
        获取当前六自由度速度
        转为机械臂控制单位mm/ms
        
        返回:
            Tuple[float, float, float, float, float, float]: (vx, vy, vz, rx, ry, rz)
        """
        return (self.vx, self.vy, self.vz, self.rx, self.ry, self.rz)
    
    def print_velocity(self):
        """打印当前速度信息"""
        print(f"线速度: X={self.vx:.3f}mm/s, Y={self.vy:.3f}mm/s, Z={self.vz:.3f}mm/s")
        print(f"角速度: RX={self.rx:.1f}m°/s, RY={self.ry:.1f}m°/s, RZ={self.rz:.1f}m°/s")
    
    def close(self):
        """关闭游戏手柄和pygame"""
        pygame.joystick.quit()
        pygame.quit()


def enable_arm(piper: C_PiperInterface_V2, timeout: int = 5) -> bool:
    """
    使能机械臂并检测使能状态
    
    参数:
        piper: 机械臂接口
        timeout: 超时时间（秒）
        
    返回:
        bool: 是否成功使能
    """
    enable_flag = False
    start_time = time.time()
    
    while not enable_flag:
        elapsed_time = time.time() - start_time
        print("--------------------")
        
        # 检查所有电机的使能状态
        enable_flag = piper.GetArmLowSpdInfoMsgs().motor_1.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_2.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_3.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_4.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_5.foc_status.driver_enable_status and \
            piper.GetArmLowSpdInfoMsgs().motor_6.foc_status.driver_enable_status
            
        print("使能状态:", enable_flag)
        piper.EnableArm(7)  # 使能所有关节
        piper.GripperCtrl(0, 1000, 0x01, 0)  # 使能夹爪
        print("--------------------")
        
        # 检查是否超时
        if elapsed_time > timeout:
            print("使能超时...")
            return False
            
        time.sleep(1)
    
    return True


def print_instructions():
    """打印操作说明"""
    print("\n==== 游戏手柄控制机械臂 ====")
    print("左摇杆: 控制X和Y轴移动")
    print("右摇杆: 控制Z轴移动和RZ轴旋转")
    print("左触发键(LT): 增加RX轴正向旋转")
    print("右触发键(RT): 增加RX轴负向旋转")
    print("LB/RB: 增加/减少速度因子")
    print("Start: 重置所有速度为零")
    print("Back/Select: 紧急停止")
    print("==========================\n")


def main():
    """主函数"""
    # 初始化机械臂接口
    piper = C_PiperInterface_V2("can0")
    piper.ConnectPort()
    piper.EnableArm(7)
    # 使能夹爪
    piper.GripperCtrl(0,1000,0x01, 0)
    
    # 使能机械臂
    print("正在使能机械臂...")
    if not enable_arm(piper):
        print("机械臂使能失败，退出程序")
        return
    

    
    # 获取当前末端位姿
    end_pose = piper.GetArmEndPoseMsgs()
    current_pose = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    
    if end_pose:
        current_pose[0] = end_pose.end_pose.X_axis / factor
        current_pose[1] = end_pose.end_pose.Y_axis / factor
        current_pose[2] = end_pose.end_pose.Z_axis / factor
        current_pose[3] = end_pose.end_pose.RX_axis / factor
        current_pose[4] = end_pose.end_pose.RY_axis / factor
        current_pose[5] = end_pose.end_pose.RZ_axis / factor
        
        print(f"当前位姿: X={current_pose[0]:.3f}mm, Y={current_pose[1]:.3f}mm, Z={current_pose[2]:.3f}mm, " +
              f"RX={current_pose[3]:.1f}°, RY={current_pose[4]:.1f}°, RZ={current_pose[5]:.1f}°")
    
    # 设置控制模式为位置速度模式
    piper.MotionCtrl_2(0x01, 0x00, 50, 0x00)
    
    try:
        # 初始化游戏手柄控制器
        controller = JoystickController(piper=piper)
        
        # 打印操作说明
        print_instructions()
        
        # 主循环
        running = True
        last_update_time = time.time()
        
        while running:

            end_pose = piper.GetArmEndPoseMsgs()
            if end_pose:
                current_pose[0] = end_pose.end_pose.X_axis / factor #mm
                current_pose[1] = end_pose.end_pose.Y_axis / factor #mm
                current_pose[2] = end_pose.end_pose.Z_axis / factor #mm
                current_pose[3] = end_pose.end_pose.RX_axis / factor #°
                current_pose[4] = end_pose.end_pose.RY_axis / factor #°
                current_pose[5] = end_pose.end_pose.RZ_axis / factor #°
                
            # 处理游戏手柄事件
            running = controller.process_events()
            
            # 更新速度
            controller.update_velocity()
            
            # 获取六自由度速度
            vx, vy, vz, rx, ry, rz = controller.get_velocity()
            
            
            # 计算时间增量
            current_time = time.time()
            dt = current_time - last_update_time
            last_update_time = current_time
            
            # 根据速度更新位置
            current_pose[0] += vx 
            current_pose[1] += vy 
            current_pose[2] += vz 
            current_pose[3] += rx 
            current_pose[4] += ry 
            current_pose[5] += rz       
            X = round(current_pose[0] * factor)                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                           
            Y = round(current_pose[1] * factor)
            Z = round(current_pose[2] * factor)
            RX = round(current_pose[3] * factor)
            RY = round(current_pose[4] * factor)
            RZ = round(current_pose[5] * factor)
            
            # 设置机械臂末端位姿
            piper.MotionCtrl_2(0x01, 0x00, 100, 0x00)
            piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            
            # 每秒打印一次当前速度信息
            if int(current_time) != int(current_time - dt):
                controller.print_velocity()
                print(f"当前位姿: X={current_pose[0]:.3f}mm, Y={current_pose[1]:.3f}mm, Z={current_pose[2]:.3f}mm, " +
                      f"RX={current_pose[3]:.1f}m°, RY={current_pose[4]:.1f}m°, RZ={current_pose[5]:.1f}m°")
            
            # 短暂延时，避免过于频繁的控制命令
            time.sleep(0.01)
    
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        # 关闭游戏手柄
        if 'controller' in locals():
            controller.close()
        
        # 停止机械臂
        # piper.MotionCtrl_1(0x02, 0, 0)  # 恢复
        # piper.MotionCtrl_2(0, 0, 0, 0x00)  # 位置速度模式
        print("程序已退出")


if __name__ == "__main__":
    main()
